import time


def torch2onnx():
    from torch2trt.torch2onnx import Torch2onnx, Config
    opt = Config()
    t2o = Torch2onnx(opt)
    t2o.convert()


def onnx2trt():
    from torch2trt.onnx2trt import Config, TrtModel
    opt = Config()
    trt = TrtModel(opt)
    trt.build_engine()


def testrtr():
    from torch2trt.onnx2trt import Config, TrtModel
    opt = Config()
    trtmodel = TrtModel(opt)
    trtmodel.load_from_trt()
    titles = ['英国首相府违规聚会调查报告认为政府“领导不力”---调查报告认为政府“领导不力”',
              '在欢声笑语中展现新时代新征程上精气神---2022年春节联欢晚会——在欢声笑语中展现新时代新征程上精气神',
              '因发生斗殴事件致2名犯人死亡 美国联邦监狱进入封锁状态---美国休斯敦发生枪击事件 致1名警员死亡',
              '英媒：泽连斯基下令三年扩军十万 敦促议员不要散布恐慌---泽连斯基签署法令：未来三年内扩军十万人',
              '冬季风暴渐平息 美部分地区降雪超50厘米---冬季风暴持续影响美国大部分地区',
              '英媒：特朗普大举筹款瞄准2024总统大选---特朗普又夸下海口：若2024再当选总统，将赦免国会大厦骚乱者',
              '多国政要和国际组织官员贺新春 祝新年如虎添翼---视频｜多位国际组织负责人及国家政要送上新春祝福',
              '“美国‘超额死亡’人数已近百万”---华尔街日报：疫情下，美国“超额死亡”人数已近百万',
              '日本外相：驻日美军入境新冠检测所用方法未被日本认可有效---外媒：驻日美军人员被曝离美前未进行新冠检测',
              '除夕，布林肯又发新闻公报拜年：愿虎年给所有人带来机遇、成功和健康---美国务卿布林肯拜年：愿虎年给所有人带来机遇、成功和健康']
    for i in titles:
        ori, sim = i.split('---')
        print(trtmodel(ori, sim))


if __name__ == '__main__':
    torch2onnx()
    # testrtr()
    # from inference import ptConfig, PtModel
    # from tqdm import tqdm
    # from torch2trt.onnx2trt import Config, TrtModel
    # import numpy as np
    # from numpy import mean
    #
    # opt = ptConfig()
    # ptm = PtModel(opt)
    #
    # data = open('testdata/devres.txt', 'r', encoding='utf-8').readlines()
    # resd = open('ptmodel--trtmodel.txt', 'a+', encoding='utf-8')
    # opt1 = Config()
    # trtmodel = TrtModel(opt1)
    # trtmodel.load_from_trt()
    # orititlelist = []
    # simtitlelist = []
    # score = []
    # timecost = []
    # for i in tqdm(data[:100], desc='ptmodel'):
    #     line = i.strip().split('\t')
    #     if len(line) == 7:
    #         ori, sim = line[3], line[4]
    #         t1 = time.time()
    #         res = ptm.inference(ori, sim)
    #         t2 = time.time()
    #         cost = (t2 - t1) * 1000
    #         resd.write(f'ptmodel||index:{line[0]}--simscore:{res.get("simscore")}--cost{cost}ms\n')
    #         orititlelist.append(res.get('orititle'))
    #         simtitlelist.append(res.get('simtitle'))
    #         score.append(res.get('simscore'))
    #         timecost.append(cost)
    # orititlelist1 = []
    # simtitlelist1 = []
    # score1 = []
    # timecost1 = []
    # resd.write(f'\r\n' * 5)
    #
    # for i in tqdm(data[:100], desc='trtmodel'):
    #     line = i.strip().split('\t')
    #     if len(line) == 7:
    #         ori, sim = line[3], line[4]
    #         t1 = time.time()
    #         res = trtmodel(ori, sim)
    #         t2 = time.time()
    #         cost = (t2 - t1) * 1000
    #         resd.write(f'trtmodel||index:{line[0]}--simscore:{res.get("simscore")}--cost{cost}ms\n')
    #         orititlelist1.append(res.get('orititle'))
    #         simtitlelist1.append(res.get('simtitle'))
    #         score1.append(res.get('simscore'))
    #         timecost1.append(cost)
    # resd.write(f'\r\n' * 5)
    # avdisscore = []
    # avdiscost = []
    # for i in range(len(timecost)):
    #     ori = orititlelist[i] - orititlelist1[i]
    #     sim = simtitlelist[i] - simtitlelist1[i]
    #     disscore = score[i] - score1[i]
    #     discost = timecost[i] - timecost1[i]
    #     print(disscore, discost)
    #     avdisscore.append(disscore)
    #     avdiscost.append(discost)
    #     resd.write(f'diff||disscore:{str(disscore)}--discost:{discost}\n')
    # resd.write(f'\r\n' * 5)
    # resd.write(f'avg diff score: {str(mean(avdisscore))}\n')
    # resd.write(f'avg diff cost: {str(mean(avdiscost))}\n')
    # resd.write(f'ptmodel avg cost: {str(mean(timecost))}\n')
    # resd.write(f'trtmodel avg cost: {str(mean(timecost1))}\n')
